import asyncio
import time
from typing import Optional, Dict
from loguru import logger
from utils.tll_consul import ConsulClient

class ConsulCleanupTask:
    """Consul服务清理后台任务 - 带失效计数机制"""
    
    def __init__(self, 
                 cleanup_interval: int = 120,  # 2分钟检查一次
                 max_critical_age: int = 10,   # critical状态超过10分钟则清理
                 fail_count_threshold: int = 2): # 失效2次后才清理
        self.cleanup_interval = cleanup_interval
        self.max_critical_age = max_critical_age
        self.fail_count_threshold = fail_count_threshold
        self.consul_client = ConsulClient()
        self._cleanup_task: Optional[asyncio.Task] = None
        self._running = False
        
        # 跟踪失效服务实例的计数
        self._failed_services: Dict[str, int] = {}  # service_id -> fail_count
        self._last_check_time = time.time()
    
    async def start_cleanup_task(self):
        """启动定期清理任务"""
        if self._running:
            logger.warning("清理任务已经在运行中")
            return
            
        self._running = True
        self._cleanup_task = asyncio.create_task(self._cleanup_loop())
        logger.info(f"Consul清理任务已启动，检查间隔: {self.cleanup_interval}秒，失效{self.fail_count_threshold}次后清理")
    
    async def stop_cleanup_task(self):
        """停止清理任务"""
        if not self._running:
            return
            
        self._running = False
        if self._cleanup_task:
            self._cleanup_task.cancel()
            try:
                await self._cleanup_task
            except asyncio.CancelledError:
                pass
        logger.info("Consul清理任务已停止")
    
    async def _cleanup_loop(self):
        """清理任务循环"""
        while self._running:
            try:
                await self._perform_cleanup()
                await asyncio.sleep(self.cleanup_interval)
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"清理任务执行异常: {e}")
                await asyncio.sleep(30)  # 出错后等待30秒再重试
    
    async def _perform_cleanup(self):
        """执行清理操作 - 带失效计数机制"""
        try:
            current_time = time.time()
            
            # 获取所有critical状态的服务
            critical_services = self.consul_client.get_critical_services()
            
            # 当前critical服务的ID集合
            current_critical_ids = set()
            
            if critical_services:
                for service in critical_services:
                    service_id = service.get('service_id')
                    if service_id:
                        current_critical_ids.add(service_id)
                        
                        # 增加或更新失效计数
                        if service_id in self._failed_services:
                            self._failed_services[service_id] += 1
                        else:
                            self._failed_services[service_id] = 1
                            logger.info(f"检测到新的critical服务: {service.get('service_name', 'unknown')} (ID: {service_id}), 开始计数")
            
            # 清理已经恢复正常的服务记录
            recovered_services = set(self._failed_services.keys()) - current_critical_ids
            for service_id in recovered_services:
                logger.info(f"服务已恢复正常, 清除失效记录: {service_id}")
                del self._failed_services[service_id]
            
            # 检查需要清理的服务（失效次数达到阈值）
            services_to_cleanup = []
            for service_id, fail_count in self._failed_services.items():
                if fail_count >= self.fail_count_threshold:
                    services_to_cleanup.append(service_id)
            
            # 执行清理
            cleaned_count = 0
            for service_id in services_to_cleanup:
                try:
                    # 查找服务信息
                    service_info = None
                    if critical_services:
                        for service in critical_services:
                            if service.get('service_id') == service_id:
                                service_info = service
                                break
                    
                    if service_info:
                        # 执行清理
                        result = self.consul_client.client.agent.service.deregister(service_id)
                        logger.info(f"清理critical服务: {service_info.get('service_name', 'unknown')} (ID: {service_id}), 失效次数: {self._failed_services[service_id]}")
                        cleaned_count += 1
                        
                        # 从跟踪列表中移除
                        del self._failed_services[service_id]
                        
                except Exception as e:
                    logger.error(f"清理服务 {service_id} 失败: {e}")
            
            if cleaned_count > 0:
                logger.info(f"定期清理: 清理了 {cleaned_count} 个critical状态的服务实例")
            
            # 输出当前状态统计
            if self._failed_services:
                logger.debug(f"当前跟踪的失效服务: {len(self._failed_services)} 个")
                for service_id, count in self._failed_services.items():
                    logger.debug(f"  - {service_id}: 失效{count}次")
            
            self._last_check_time = current_time
            
        except Exception as e:
            logger.error(f"执行清理操作失败: {e}")
    
    def get_failed_services_status(self) -> Dict[str, int]:
        """获取当前失效服务的状态"""
        return self._failed_services.copy()
    
    def reset_service_fail_count(self, service_id: str) -> bool:
        """重置指定服务的失效计数"""
        if service_id in self._failed_services:
            del self._failed_services[service_id]
            logger.info(f"已重置服务 {service_id} 的失效计数")
            return True
        return False
    
    def reset_all_fail_counts(self):
        """重置所有服务的失效计数"""
        count = len(self._failed_services)
        self._failed_services.clear()
        logger.info(f"已重置所有 {count} 个服务的失效计数")
        return count

# 全局清理任务实例
cleanup_task = ConsulCleanupTask()

async def start_consul_cleanup():
    """启动Consul清理任务（在应用启动时调用）"""
    await cleanup_task.start_cleanup_task()

async def stop_consul_cleanup():
    """停止Consul清理任务（在应用关闭时调用）"""
    await cleanup_task.stop_cleanup_task()